package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

type Resource struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

type MCPServer struct {
	db *sql.DB
}

func NewMCPServer() (*MCPServer, error) {
	db, err := sql.Open("mysql", "root:1AVPeG4ofQmMYaruGO88@tcp(tt-sit-rds-mysql.xtt-inc.xyz:3306)/wtest?parseTime=true&charset=utf8&loc=Local")
	if err != nil {
		return nil, err
	}
	return &MCPServer{db: db}, nil
}

func (s *MCPServer) handleConnection(conn *websocket.Conn) {
	defer conn.Close()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println("Read error:", err)
			break
		}

		var request map[string]interface{}
		if err := json.Unmarshal(msg, &request); err != nil {
			sendError(conn, "Invalid request format")
			continue
		}

		switch request["action"] {
		case "get_resource":
			s.handleGetResource(conn, request)
		default:
			sendError(conn, "Unsupported action")
		}
	}
}

func (s *MCPServer) handleGetResource(conn *websocket.Conn, request map[string]interface{}) {
	id, ok := request["resource_id"].(float64)
	if !ok {
		sendError(conn, "Invalid resource ID")
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	row := s.db.QueryRowContext(ctx, "SELECT * FROM resources WHERE id = ?", int(id))
	var res Resource
	if err := row.Scan(&res.ID, &res.Name); err != nil {
		sendError(conn, "Resource not found")
		return
	}

	response := map[string]interface{}{
		"status": "success",
		"data":   res,
	}
	conn.WriteJSON(response)
}

func sendError(conn *websocket.Conn, message string) {
	conn.WriteJSON(map[string]interface{}{
		"status":  "error",
		"message": message,
	})
}

func main() {
	server, err := NewMCPServer()
	if err != nil {
		log.Fatal("Database connection failed:", err)
	}
	defer server.db.Close()

	http.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println("WebSocket upgrade failed:", err)
			return
		}
		server.handleConnection(conn)
	})

	listener, err := net.Listen("tcp", ":8080")
	if err != nil {
		log.Fatal("Listen failed:", err)
	}

	log.Println("MCP Server running on ws://localhost:8080/mcp")
	log.Fatal(http.Serve(listener, nil))
}
